/**
 * 
 */
package edu.gatech.cc.liam.core;

import java.io.Serializable;
import java.util.Arrays;
import java.util.Map;
import java.util.Random;
import java.util.Set;
import java.util.TreeMap;

public class DiscreteDistribution<T> extends TreeMap<T,Double> implements Serializable{
	private static final long serialVersionUID = 2914587422027055774L;
	
	public T sample() {
		double target = Math.random();
		for (Map.Entry<T,Double> e : this.entrySet()) {
			target -= e.getValue();
			if (target < 0)
				return e.getKey();
		}
		return null;
	}
	
	@Override
	public Double put(T key, Double value) {
		Double previousValue = super.put(key, value);
		if(previousValue != null) {
			super.put(key, previousValue + value);
		}
		return previousValue;
	}

	public Set<Map.Entry<T,Double>> getEntrySet() {
		return this.entrySet();
	}
	
	public boolean isDistribution() {
		double sum = 0.0;
		for (Map.Entry<T,Double> e : this.entrySet()) {
			sum += e.getValue();
		}
		return sum == 1.0;
	}
	
	public void normalize() {
		double sum = 0.0;
		for (Map.Entry<T,Double> e : this.entrySet()) {
			sum += e.getValue();
		}
		for (Map.Entry<T,Double> e : this.entrySet()) {
			super.put(e.getKey(), e.getValue() / sum);
		}
	}
	
	public static int sampleDist(double[] distribution) {
		double target = Math.random();
		double curVal = 0.0;
		for (int i=0; i<distribution.length; i++) {
			curVal += distribution[i];
			if (target < curVal)
				return i;
		}
		return -1;
	}
	
	/**
	 * correctly samples from the uniform distribution across a simplex
	 * @param numDimentions
	 * @param rand
	 * @return
	 */
	public static double[] getRandomDistribution(int numDimentions, Random rand) {
		double[] randValues = new double[numDimentions+1];
		for(int i=0; i<numDimentions-1; i++) {
			randValues[i] = rand.nextDouble();
		}
		randValues[numDimentions-1] = 0.0;
		randValues[numDimentions] = 1.0;
		Arrays.sort(randValues);
		double[] rVals = new double[numDimentions];
		for(int i=0; i<numDimentions; i++) {
			rVals[i] = randValues[i+1] - randValues[i];
		}
		return rVals;
	}
}
